import os

from dotenv import load_dotenv
from langchain.chat_models import init_chat_model
from langchain_core.messages import AIMessage, HumanMessage, SystemMessage
from langchain_core.output_parsers import StrOutputParser
from langchain_core.prompts import ChatPromptTemplate, MessagesPlaceholder

if __name__ == '__main__':
    load_dotenv(override=True)

    chatbot_prompt = ChatPromptTemplate.from_messages(
        [
            SystemMessage(content="你叫做小智，是一个乐于助人的助手"),
            MessagesPlaceholder(variable_name="messages"),
        ]
    )
    DEEPSEEK_API_KEY = os.getenv("DEEPSEEK_API_KEY")
    model = init_chat_model(model="deepseek-chat", model_provider="deepseek")
    basic_qa_chain = chatbot_prompt | model | StrOutputParser()

    message_list = [
        HumanMessage(content="你好，我叫做陈明，好久不见"),
        AIMessage(content="你好，陈明，很高兴见到你,你叫做小智，是一个乐于助人的助手"),
    ]
    question = "你好，请问我叫做什么名字"
    message_list.append(HumanMessage(content=question))
    print(message_list)
    result = basic_qa_chain.invoke({"messages": message_list})
    print(result)